kreuzberg 3.7.0__py3-none-any.whl → 3.8.1__py3-none-any.whl

kreuzberg/_entity_extraction.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 import re
 from dataclasses import dataclass
 from functools import lru_cache
@@ -181,8 +182,6 @@ def _load_spacy_model(model_name: str, spacy_config: SpacyEntityExtractionConfig
         import spacy
 
         if spacy_config.model_cache_dir:
-            import os
-
             os.environ["SPACY_DATA"] = str(spacy_config.model_cache_dir)
 
         nlp = spacy.load(model_name)

kreuzberg/_extractors/_base.py

@@ -3,10 +3,12 @@ from __future__ import annotations
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, ClassVar
 
+from kreuzberg._types import ExtractionResult, normalize_metadata
+from kreuzberg._utils._quality import calculate_quality_score, clean_extracted_text
+
 if TYPE_CHECKING:
     from pathlib import Path
 
-    from kreuzberg import ExtractionResult
     from kreuzberg._types import ExtractionConfig
 
 
@@ -90,3 +92,39 @@ class Extractor(ABC):
         return mime_type in cls.SUPPORTED_MIME_TYPES or any(
             mime_type.startswith(supported_type) for supported_type in cls.SUPPORTED_MIME_TYPES
         )
+
+    def _apply_quality_processing(self, result: ExtractionResult) -> ExtractionResult:
+        """Apply quality post-processing to extraction result if enabled.
+
+        Args:
+            result: The raw extraction result
+
+        Returns:
+            Enhanced extraction result with quality improvements (if enabled)
+        """
+        # Only apply quality processing if enabled in config
+        if not self.config.enable_quality_processing:
+            return result
+
+        if not result.content:
+            return result
+
+        # Clean the content
+        cleaned_content = clean_extracted_text(result.content)
+
+        # Calculate quality score
+        quality_score = calculate_quality_score(cleaned_content, dict(result.metadata) if result.metadata else None)
+
+        # Add quality metadata
+        enhanced_metadata = dict(result.metadata) if result.metadata else {}
+        enhanced_metadata["quality_score"] = quality_score
+
+        # Return enhanced result
+        return ExtractionResult(
+            content=cleaned_content,
+            mime_type=result.mime_type,
+            metadata=normalize_metadata(enhanced_metadata),
+            chunks=result.chunks,
+            detected_languages=result.detected_languages,
+            tables=result.tables,
+        )

kreuzberg/_extractors/_email.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+import re
+from html import unescape
+from typing import TYPE_CHECKING, Any, ClassVar
+
+from anyio import Path as AsyncPath
+
+from kreuzberg._extractors._base import Extractor
+from kreuzberg._mime_types import EML_MIME_TYPE, PLAIN_TEXT_MIME_TYPE
+from kreuzberg._types import ExtractionResult, normalize_metadata
+from kreuzberg._utils._string import normalize_spaces
+from kreuzberg._utils._sync import run_sync
+from kreuzberg.exceptions import MissingDependencyError
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+# Import optional dependencies at module level with proper error handling
+try:
+    import mailparse
+except ImportError:
+    mailparse = None
+
+try:
+    import html2text  # type: ignore[import-not-found]
+except ImportError:
+    html2text = None
+
+# Compile regex pattern once at module level
+_HTML_TAG_PATTERN = re.compile(r"<[^>]+>")
+
+
+class EmailExtractor(Extractor):
+    SUPPORTED_MIME_TYPES: ClassVar[set[str]] = {EML_MIME_TYPE}
+
+    async def extract_bytes_async(self, content: bytes) -> ExtractionResult:
+        return await run_sync(self.extract_bytes_sync, content)
+
+    async def extract_path_async(self, path: Path) -> ExtractionResult:
+        content = await AsyncPath(path).read_bytes()
+        return await self.extract_bytes_async(content)
+
+    def _extract_email_headers(
+        self, parsed_email: dict[str, Any], text_parts: list[str], metadata: dict[str, Any]
+    ) -> None:
+        """Extract and process email headers."""
+        # Use single dict access where possible to avoid repeated lookups
+        subject = parsed_email.get("subject")
+        if subject:
+            metadata["subject"] = subject
+            text_parts.append(f"Subject: {subject}")
+
+        from_info = parsed_email.get("from")
+        if from_info:
+            from_email = from_info.get("email", "") if isinstance(from_info, dict) else str(from_info)
+            metadata["email_from"] = from_email
+            text_parts.append(f"From: {from_email}")
+
+        to_info = parsed_email.get("to")
+        if to_info:
+            if isinstance(to_info, list) and to_info:
+                to_email = to_info[0].get("email", "") if isinstance(to_info[0], dict) else str(to_info[0])
+            elif isinstance(to_info, dict):
+                to_email = to_info.get("email", "")
+            else:
+                to_email = str(to_info)
+            metadata["email_to"] = to_email
+            text_parts.append(f"To: {to_email}")
+
+        date = parsed_email.get("date")
+        if date:
+            metadata["date"] = date
+            text_parts.append(f"Date: {date}")
+
+        cc = parsed_email.get("cc")
+        if cc:
+            metadata["email_cc"] = cc
+            text_parts.append(f"CC: {cc}")
+
+        bcc = parsed_email.get("bcc")
+        if bcc:
+            metadata["email_bcc"] = bcc
+            text_parts.append(f"BCC: {bcc}")
+
+    def _extract_email_body(self, parsed_email: dict[str, Any], text_parts: list[str]) -> None:
+        """Extract and process email body content."""
+        text_content = parsed_email.get("text")
+        if text_content:
+            text_parts.append(f"\n{text_content}")
+            return  # If we have text, prefer it over HTML
+
+        html_content = parsed_email.get("html")
+        if html_content:
+            if html2text is not None:
+                # Use html2text if available (faster path)
+                h = html2text.HTML2Text()
+                h.ignore_links = True
+                h.ignore_images = True
+                converted_text = h.handle(html_content)
+                text_parts.append(f"\n{converted_text}")
+            else:
+                # Fallback: strip HTML tags and unescape entities
+                clean_html = _HTML_TAG_PATTERN.sub("", html_content)
+                clean_html = unescape(clean_html)
+                text_parts.append(f"\n{clean_html}")
+
+    def _extract_email_attachments(
+        self, parsed_email: dict[str, Any], text_parts: list[str], metadata: dict[str, Any]
+    ) -> None:
+        """Extract and process email attachments info."""
+        if parsed_email.get("attachments"):
+            attachment_names = [att.get("name", "unknown") for att in parsed_email["attachments"]]
+            metadata["attachments"] = attachment_names
+            if attachment_names:
+                text_parts.append(f"\nAttachments: {', '.join(attachment_names)}")
+
+    def extract_bytes_sync(self, content: bytes) -> ExtractionResult:
+        if mailparse is None:
+            msg = "mailparse is required for email extraction. Install with: pip install 'kreuzberg[additional-extensions]'"
+            raise MissingDependencyError(msg)
+
+        try:
+            parsed_email = mailparse.EmailDecode.load(content)
+            text_parts: list[str] = []
+            metadata: dict[str, Any] = {}
+
+            # Extract headers, body, and attachments
+            self._extract_email_headers(parsed_email, text_parts, metadata)
+            self._extract_email_body(parsed_email, text_parts)
+            self._extract_email_attachments(parsed_email, text_parts, metadata)
+
+            # Join efficiently
+            combined_text = "\n".join(text_parts)
+
+            return ExtractionResult(
+                content=normalize_spaces(combined_text),
+                mime_type=PLAIN_TEXT_MIME_TYPE,
+                metadata=normalize_metadata(metadata),
+                chunks=[],
+            )
+
+        except Exception as e:
+            msg = f"Failed to parse email content: {e}"
+            raise RuntimeError(msg) from e
+
+    def extract_path_sync(self, path: Path) -> ExtractionResult:
+        content = path.read_bytes()
+        return self.extract_bytes_sync(content)

kreuzberg/_extractors/_html.py

@@ -8,7 +8,7 @@ from anyio import Path as AsyncPath
 from kreuzberg._extractors._base import Extractor
 from kreuzberg._mime_types import HTML_MIME_TYPE, MARKDOWN_MIME_TYPE
 from kreuzberg._types import ExtractionResult
-from kreuzberg._utils._string import normalize_spaces, safe_decode
+from kreuzberg._utils._string import safe_decode
 from kreuzberg._utils._sync import run_sync
 
 if TYPE_CHECKING:
@@ -26,8 +26,20 @@ class HTMLExtractor(Extractor):
         return await run_sync(self.extract_bytes_sync, content)
 
     def extract_bytes_sync(self, content: bytes) -> ExtractionResult:
-        result = html_to_markdown.convert_to_markdown(safe_decode(content))
-        return ExtractionResult(content=normalize_spaces(result), mime_type=MARKDOWN_MIME_TYPE, metadata={}, chunks=[])
+        # Use html-to-markdown with script/nav removal for better quality
+        result = html_to_markdown.convert_to_markdown(
+            safe_decode(content),
+            preprocess_html=True,
+            preprocessing_preset="aggressive",
+            remove_navigation=True,
+            remove_forms=True,
+        )
+
+        # Skip normalize_spaces since quality processing will handle whitespace
+        extraction_result = ExtractionResult(content=result, mime_type=MARKDOWN_MIME_TYPE, metadata={}, chunks=[])
+
+        # Apply quality processing which includes normalization
+        return self._apply_quality_processing(extraction_result)
 
     def extract_path_sync(self, path: Path) -> ExtractionResult:
         content = path.read_bytes()

kreuzberg/_extractors/_image.py

@@ -1,5 +1,9 @@
 from __future__ import annotations
 
+import contextlib
+import os
+import tempfile
+from pathlib import Path
 from typing import TYPE_CHECKING, ClassVar
 
 from anyio import Path as AsyncPath
@@ -7,6 +11,9 @@ from anyio import Path as AsyncPath
 from kreuzberg._extractors._base import Extractor
 from kreuzberg._mime_types import IMAGE_MIME_TYPES
 from kreuzberg._ocr import get_ocr_backend
+from kreuzberg._ocr._easyocr import EasyOCRConfig
+from kreuzberg._ocr._paddleocr import PaddleOCRConfig
+from kreuzberg._ocr._tesseract import TesseractConfig
 from kreuzberg._utils._tmp import create_temp_file
 from kreuzberg.exceptions import ValidationError
 
@@ -15,9 +22,6 @@ if TYPE_CHECKING:  # pragma: no cover
 
     from kreuzberg._types import ExtractionResult
 
-import contextlib
-from pathlib import Path
-
 
 class ImageExtractor(Extractor):
     SUPPORTED_MIME_TYPES: ClassVar[set[str]] = IMAGE_MIME_TYPES
@@ -56,13 +60,11 @@ class ImageExtractor(Extractor):
         if self.config.ocr_backend is None:
             raise ValidationError("ocr_backend is None, cannot perform OCR")
 
-        return await get_ocr_backend(self.config.ocr_backend).process_file(path, **self.config.get_config_dict())
+        result = await get_ocr_backend(self.config.ocr_backend).process_file(path, **self.config.get_config_dict())
+        return self._apply_quality_processing(result)
 
     def extract_bytes_sync(self, content: bytes) -> ExtractionResult:
         """Pure sync implementation of extract_bytes."""
-        import os
-        import tempfile
-
         extension = self._get_extension_from_mime_type(self.mime_type)
         fd, temp_path = tempfile.mkstemp(suffix=f".{extension}")
 
@@ -80,43 +82,26 @@ class ImageExtractor(Extractor):
         if self.config.ocr_backend is None:
             raise ValidationError("ocr_backend is None, cannot perform OCR")
 
-        from kreuzberg._types import ExtractionResult
+        backend = get_ocr_backend(self.config.ocr_backend)
 
         if self.config.ocr_backend == "tesseract":
-            from kreuzberg._multiprocessing.sync_tesseract import process_batch_images_sync_pure
-            from kreuzberg._ocr._tesseract import TesseractConfig
-
-            if isinstance(self.config.ocr_config, TesseractConfig):
-                config = self.config.ocr_config
-            else:
-                config = TesseractConfig()
-
-            results = process_batch_images_sync_pure([str(path)], config)
-            if results:
-                return results[0]
-            return ExtractionResult(content="", mime_type="text/plain", metadata={}, chunks=[])
-
-        if self.config.ocr_backend == "paddleocr":
-            from kreuzberg._multiprocessing.sync_paddleocr import process_image_sync_pure as paddle_process
-            from kreuzberg._ocr._paddleocr import PaddleOCRConfig
-
+            config = (
+                self.config.ocr_config if isinstance(self.config.ocr_config, TesseractConfig) else TesseractConfig()
+            )
+            result = backend.process_file_sync(path, **config.__dict__)
+        elif self.config.ocr_backend == "paddleocr":
             paddle_config = (
                 self.config.ocr_config if isinstance(self.config.ocr_config, PaddleOCRConfig) else PaddleOCRConfig()
             )
-
-            return paddle_process(path, paddle_config)
-
-        if self.config.ocr_backend == "easyocr":
-            from kreuzberg._multiprocessing.sync_easyocr import process_image_sync_pure as easy_process
-            from kreuzberg._ocr._easyocr import EasyOCRConfig
-
+            result = backend.process_file_sync(path, **paddle_config.__dict__)
+        elif self.config.ocr_backend == "easyocr":
             easy_config = (
                 self.config.ocr_config if isinstance(self.config.ocr_config, EasyOCRConfig) else EasyOCRConfig()
             )
-
-            return easy_process(path, easy_config)
-
-        raise NotImplementedError(f"Sync OCR not implemented for {self.config.ocr_backend}")
+            result = backend.process_file_sync(path, **easy_config.__dict__)
+        else:
+            raise NotImplementedError(f"Sync OCR not implemented for {self.config.ocr_backend}")
+        return self._apply_quality_processing(result)
 
     def _get_extension_from_mime_type(self, mime_type: str) -> str:
         if mime_type in self.IMAGE_MIME_TYPE_EXT_MAP:

kreuzberg/_extractors/_pandoc.py

@@ -1,8 +1,11 @@
 from __future__ import annotations
 
 import contextlib
+import os
 import re
+import subprocess
 import sys
+import tempfile
 from json import JSONDecodeError, loads
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast
@@ -203,10 +206,6 @@ class PandocExtractor(Extractor):
         Returns:
             ExtractionResult with the extracted text and metadata.
         """
-        import os
-        import tempfile
-        from pathlib import Path
-
         extension = self._get_pandoc_type_from_mime_type(self.mime_type)
         fd, temp_path = tempfile.mkstemp(suffix=f".{extension}")
 
@@ -579,8 +578,6 @@ class PandocExtractor(Extractor):
 
     def _validate_pandoc_version_sync(self) -> None:
         """Synchronous version of _validate_pandoc_version."""
-        import subprocess
-
         try:
             if self._checked_version:
                 return
@@ -625,10 +622,6 @@ class PandocExtractor(Extractor):
 
     def _extract_metadata_sync(self, path: Path) -> Metadata:
         """Synchronous version of _handle_extract_metadata."""
-        import os
-        import subprocess
-        import tempfile
-
         pandoc_type = self._get_pandoc_type_from_mime_type(self.mime_type)
         fd, metadata_file = tempfile.mkstemp(suffix=".json")
         os.close(fd)
@@ -663,10 +656,6 @@ class PandocExtractor(Extractor):
 
     def _extract_file_sync(self, path: Path) -> str:
         """Synchronous version of _handle_extract_file."""
-        import os
-        import subprocess
-        import tempfile
-
         pandoc_type = self._get_pandoc_type_from_mime_type(self.mime_type)
         fd, output_path = tempfile.mkstemp(suffix=".md")
         os.close(fd)

kreuzberg/_extractors/_pdf.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import contextlib
+import os
+import tempfile
 from multiprocessing import cpu_count
 from pathlib import Path
 from re import Pattern
@@ -10,15 +12,21 @@ from typing import TYPE_CHECKING, ClassVar, cast
 import anyio
 import pypdfium2
 from anyio import Path as AsyncPath
+from playa import parse
 
 from kreuzberg._extractors._base import Extractor
 from kreuzberg._mime_types import PDF_MIME_TYPE, PLAIN_TEXT_MIME_TYPE
 from kreuzberg._ocr import get_ocr_backend
-from kreuzberg._playa import extract_pdf_metadata
+from kreuzberg._ocr._easyocr import EasyOCRConfig
+from kreuzberg._ocr._paddleocr import PaddleOCRConfig
+from kreuzberg._ocr._tesseract import TesseractConfig
+from kreuzberg._playa import extract_pdf_metadata, extract_pdf_metadata_sync
 from kreuzberg._types import ExtractionResult, OcrBackendType
+from kreuzberg._utils._errors import create_error_context, should_retry
 from kreuzberg._utils._pdf_lock import pypdfium_file_lock
 from kreuzberg._utils._string import normalize_spaces
 from kreuzberg._utils._sync import run_sync, run_taskgroup_batched
+from kreuzberg._utils._table import generate_table_summary
 from kreuzberg._utils._tmp import create_temp_file
 from kreuzberg.exceptions import ParsingError
 
@@ -63,17 +71,30 @@ class PDFExtractor(Extractor):
         result.metadata = await extract_pdf_metadata(content_bytes)
 
         if self.config.extract_tables:
-            from kreuzberg._gmft import extract_tables
-
-            result.tables = await extract_tables(path, self.config.gmft_config)
+            # GMFT is optional dependency
+            try:
+                from kreuzberg._gmft import extract_tables
 
-        return result
+                result.tables = await extract_tables(path, self.config.gmft_config)
+            except ImportError:
+                result.tables = []
+
+            # Enhance metadata with table information
+            if result.tables:
+                table_summary = generate_table_summary(result.tables)
+                result.metadata.update(
+                    {
+                        "table_count": table_summary["table_count"],
+                        "tables_summary": f"Document contains {table_summary['table_count']} tables "
+                        f"across {table_summary['pages_with_tables']} pages with "
+                        f"{table_summary['total_rows']} total rows",
+                    }
+                )
+
+        return self._apply_quality_processing(result)
 
     def extract_bytes_sync(self, content: bytes) -> ExtractionResult:
         """Pure sync implementation of PDF extraction from bytes."""
-        import os
-        import tempfile
-
         fd, temp_path = tempfile.mkstemp(suffix=".pdf")
         try:
             with os.fdopen(fd, "wb") as f:
@@ -81,8 +102,6 @@ class PDFExtractor(Extractor):
 
             result = self.extract_path_sync(Path(temp_path))
 
-            from kreuzberg._playa import extract_pdf_metadata_sync
-
             metadata = extract_pdf_metadata_sync(content)
             result.metadata = metadata
 
@@ -100,16 +119,21 @@ class PDFExtractor(Extractor):
 
         tables = []
         if self.config.extract_tables:
+            # GMFT is optional dependency
             try:
                 from kreuzberg._gmft import extract_tables_sync
 
                 tables = extract_tables_sync(path)
             except ImportError:
-                pass
+                tables = []
+
+        # Use playa for better text structure preservation when not using OCR
+        if not self.config.force_ocr and self._validate_extracted_text(text):
+            text = self._extract_with_playa_sync(path, fallback_text=text)
 
         text = normalize_spaces(text)
 
-        return ExtractionResult(
+        result = ExtractionResult(
             content=text,
             mime_type=PLAIN_TEXT_MIME_TYPE,
             metadata={},
@@ -117,6 +141,21 @@ class PDFExtractor(Extractor):
             chunks=[],
         )
 
+        # Enhance metadata with table information
+        if tables:
+            table_summary = generate_table_summary(tables)
+            result.metadata.update(
+                {
+                    "table_count": table_summary["table_count"],
+                    "tables_summary": f"Document contains {table_summary['table_count']} tables "
+                    f"across {table_summary['pages_with_tables']} pages with "
+                    f"{table_summary['total_rows']} total rows",
+                }
+            )
+
+        # Apply quality processing
+        return self._apply_quality_processing(result)
+
     def _validate_extracted_text(self, text: str, corruption_threshold: float = 0.05) -> bool:
         """Check if text extracted from PDF is valid or corrupted.
 
@@ -155,8 +194,6 @@ class PDFExtractor(Extractor):
         Returns:
             A list of Pillow Images.
         """
-        from kreuzberg._utils._errors import create_error_context, should_retry
-
         document: pypdfium2.PdfDocument | None = None
         last_error = None
 
@@ -228,8 +265,6 @@ class PDFExtractor(Extractor):
         Returns:
             The extracted text.
         """
-        from kreuzberg._utils._errors import create_error_context
-
         document: pypdfium2.PdfDocument | None = None
         try:
             with pypdfium_file_lock(input_file):
@@ -283,7 +318,7 @@ class PDFExtractor(Extractor):
                 text_parts = []
                 for page in pdf:
                     text_page = page.get_textpage()
-                    text = text_page.get_text_range()
+                    text = text_page.get_text_bounded()
                     text_parts.append(text)
                     text_page.close()
                     page.close()
@@ -309,9 +344,6 @@ class PDFExtractor(Extractor):
                     bitmap.close()
                     page.close()
 
-            import os
-            import tempfile
-
             image_paths = []
             temp_files = []
 
@@ -339,43 +371,44 @@ class PDFExtractor(Extractor):
 
     def _process_pdf_images_with_ocr(self, image_paths: list[str]) -> str:
         """Process PDF images with the configured OCR backend."""
-        if self.config.ocr_backend == "tesseract":
-            from kreuzberg._multiprocessing.sync_tesseract import process_batch_images_sync_pure
-            from kreuzberg._ocr._tesseract import TesseractConfig
+        backend = get_ocr_backend(self.config.ocr_backend)
+        paths = [Path(p) for p in image_paths]
 
-            tesseract_config = (
+        if self.config.ocr_backend == "tesseract":
+            config = (
                 self.config.ocr_config if isinstance(self.config.ocr_config, TesseractConfig) else TesseractConfig()
             )
-            results = process_batch_images_sync_pure([str(p) for p in image_paths], tesseract_config)
-            text_parts = [r.content for r in results]
-            return "\n\n".join(text_parts)
-
-        if self.config.ocr_backend == "paddleocr":
-            from kreuzberg._multiprocessing.sync_paddleocr import process_image_sync_pure as paddle_process
-            from kreuzberg._ocr._paddleocr import PaddleOCRConfig
-
+            results = backend.process_batch_sync(paths, **config.__dict__)
+        elif self.config.ocr_backend == "paddleocr":
             paddle_config = (
                 self.config.ocr_config if isinstance(self.config.ocr_config, PaddleOCRConfig) else PaddleOCRConfig()
             )
-
-            text_parts = []
-            for image_path in image_paths:
-                result = paddle_process(Path(image_path), paddle_config)
-                text_parts.append(result.content)
-            return "\n\n".join(text_parts)
-
-        if self.config.ocr_backend == "easyocr":
-            from kreuzberg._multiprocessing.sync_easyocr import process_image_sync_pure as easy_process
-            from kreuzberg._ocr._easyocr import EasyOCRConfig
-
+            results = backend.process_batch_sync(paths, **paddle_config.__dict__)
+        elif self.config.ocr_backend == "easyocr":
             easy_config = (
                 self.config.ocr_config if isinstance(self.config.ocr_config, EasyOCRConfig) else EasyOCRConfig()
             )
+            results = backend.process_batch_sync(paths, **easy_config.__dict__)
+        else:
+            raise NotImplementedError(f"Sync OCR not implemented for {self.config.ocr_backend}")
+
+        text_parts = [r.content for r in results]
+        return "\n\n".join(text_parts)
+
+    def _extract_with_playa_sync(self, path: Path, fallback_text: str) -> str:
+        """Extract text using playa for better structure preservation."""
+        with contextlib.suppress(Exception):
+            content = path.read_bytes()
+            document = parse(content, max_workers=1)
 
             text_parts = []
-            for image_path in image_paths:
-                result = easy_process(Path(image_path), easy_config)
-                text_parts.append(result.content)
-            return "\n\n".join(text_parts)
+            for page in document.pages:
+                # Extract text while preserving structure
+                page_text = page.extract_text()
+                if page_text and page_text.strip():
+                    text_parts.append(page_text)
+
+            if text_parts:
+                return "\n\n".join(text_parts)
 
-        raise NotImplementedError(f"Sync OCR not implemented for {self.config.ocr_backend}")
+        return fallback_text